from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.common.exceptions import TimeoutException
import time
import base64
from PIL import Image
import io
import os

def capture_scroll_screenshot_and_save_pdf(
    url, 
    output_pdf="output.pdf", 
    scroll_step=500, 
    max_scrolls=20, 
    timeout=10,
    max_page_height=1500  # 新增参数：每页最大高度（像素）
):
    """
    滚动截屏并拼接为 PDF，支持自动分页
    
    :param url: 目标网页 URL
    :param output_pdf: 输出 PDF 路径（默认：output.pdf）
    :param scroll_step: 每次滚动的像素步长（默认：500px）
    :param max_scrolls: 最大滚动次数（防止无限循环，默认：20）
    :param timeout: 内容加载超时时间（秒，默认：10）
    :param max_page_height: 每页最大高度（像素，默认：1500px，超过则自动分页）
    """
    # 配置 Chrome 选项（无头模式 + 截图优化）
    chrome_options = Options()
    chrome_options.add_argument("--headless=new")  # 无头模式（Chrome 112+）
    chrome_options.add_argument("--disable-gpu")
    chrome_options.add_argument("--no-sandbox")
    chrome_options.add_argument("--window-size=1920,1080")  # 窗口尺寸（影响截图宽度）
    chrome_options.add_argument("--hide-scrollbars")  # 隐藏滚动条

    # 初始化浏览器驱动
    driver = webdriver.Chrome(options=chrome_options)

    try:
        # 导航到目标页面并等待初始加载
        driver.get(url)
        WebDriverWait(driver, timeout).until(
            EC.presence_of_element_located(("tag name", "body"))
        )

        # 获取初始窗口尺寸
        window_height = driver.execute_script("return window.innerHeight")
        scroll_position = 0
        screenshots = []

        # 模拟滚动并截图
        for _ in range(max_scrolls):
            driver.execute_script(f"window.scrollTo(0, {scroll_position});")
            scroll_position += scroll_step
            time.sleep(0.5)  # 等待滚动稳定

            # 获取截图并解码
            screenshot_base64 = driver.get_screenshot_as_base64()
            screenshot_bytes = base64.b64decode(screenshot_base64)
            img = Image.open(io.BytesIO(screenshot_bytes))

            # 裁剪实际内容区域（去除底部空白）
            actual_height = min(window_height, img.height)
            img = img.crop((0, 0, img.width, actual_height))
            screenshots.append(img)

            # 检测是否到达页面底部
            total_height = driver.execute_script("return document.documentElement.scrollHeight")
            if scroll_position >= total_height - window_height:
                break

        if not screenshots:
            raise ValueError("未获取到任何截图，请检查页面加载或滚动逻辑")

        # 拼接所有截图为长图
        total_width = screenshots[0].width
        total_height = sum(img.height for img in screenshots)
        combined_img = Image.new("RGB", (total_width, total_height))

        y_offset = 0
        for img in screenshots:
            combined_img.paste(img, (0, y_offset))
            y_offset += img.height

        # 自动分页处理（关键新增逻辑）
        if combined_img.height > max_page_height:
            print(f"合并图片高度（{combined_img.height}px）超过阈值（{max_page_height}px），自动分页...")
            
            # 计算分页数量和分割坐标
            pages = []
            num_pages = (combined_img.height + max_page_height - 1) // max_page_height
            
            for i in range(num_pages):
                start_y = i * max_page_height
                end_y = min((i + 1) * max_page_height, combined_img.height)
                # 裁剪当前页内容（保持宽度不变）
                page_img = combined_img.crop((0, start_y, total_width, end_y))
                pages.append(page_img)
            
            # 保存为多页PDF
            pages[0].save(
                output_pdf, 
                "PDF", 
                resolution=100.0, 
                save_all=True, 
                append_images=pages[1:]
            )
            print(f"PDF 成功保存至：{output_pdf}（共 {num_pages} 页）")
        else:
            # 未超过阈值直接保存
            combined_img.save(output_pdf, "PDF", resolution=100.0)
            print(f"PDF 成功保存至：{output_pdf}（单页，高度 {combined_img.height}px）")

    except Exception as e:
        print(f"生成失败，错误信息：{str(e)}")
    finally:
        driver.quit()

# 使用示例
if __name__ == "__main__":
    target_url = "http://114.251.111.103:18080/kfs/file/read/96fff55c0628d9342c2cacbd786c0bce"
    capture_scroll_screenshot_and_save_pdf(
        url=target_url,
        output_pdf="long_page.pdf",
        scroll_step=800,
        max_scrolls=1500,
        timeout=15,
        max_page_height=15000  # 可根据需求调整（如A4纸高度约1169px@100dpi）
    )